package com.caochenlei.graph;

import java.util.LinkedList;
import java.util.PriorityQueue;
import java.util.Queue;

public class PrimMST {
    private boolean[] marked;                   //索引代表顶点，值表示当前顶点是否已经被搜索
    private Queue<Edge> mst;                    //最小生成树的各个边
    private Queue<Edge> pq;                     //优先队列用来存放横切边（包括失效的边）

    public PrimMST(EdgeWeightedGraph G) {
        this.marked = new boolean[G.size()];    //初始化marked数组
        this.mst = new LinkedList<>();          //初始化mst最小生成树
        this.pq = new PriorityQueue<>();        //初始化pq优先队列

        visit(G, 0);                        //假设G是连通的
        while (!pq.isEmpty() && (mst.size() < G.size() - 1)) {
            Edge e = pq.poll();                 //取出权重最小的那条边
            int v = e.either();                 //获取权重最小的那条边的其中一个点
            int w = e.other(v);                 //获取权重最小的那条边的另外一个点
            if (marked[v] && marked[w]) {       //跳过失效的边
                continue;
            }
            mst.add(e);                         //将权重最小的那条边加入到最小生成树中
            if (!marked[v]) {                   //将顶点v添加到最小生成树中
                visit(G, v);
            }
            if (!marked[w]) {                   //将顶点w添加到最小生成树中
                visit(G, w);
            }
        }
    }

    private void visit(EdgeWeightedGraph G, int v) {
        marked[v] = true;                       //把顶点v标识为已搜索
        for (Edge e : G.adj(v)) {
            int w = e.other(v);                 //获取边e的另外一个顶点w
            if (!marked[w]) {                   //判断当前w顶点有没有被搜索过
                pq.add(e);                      //如果没有被搜索过，就直接把顶点v和顶点w所对应的边e加入到优先队列pq中
            }
        }
    }

    //获取最小生成树的各边
    public Queue<Edge> edges() {
        return mst;
    }

    //获取最小生成树的权重
    public double weight() {
        double s = 0.0D;
        if (mst != null) {
            for (Edge e : mst) {
                s += e.weight();
            }
        }
        return s;
    }
}
